{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4f741799e60ff1ae",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Retrieval-Augmented Generation\n",
    "\n",
    "This guide shows how to apply a guardrails configuration in a RAG scenario. This guide builds on the [previous guide](../6-topical-rails), developing further the demo ABC Bot.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f11740de9875c6f9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-06T20:32:41.670537Z",
     "start_time": "2023-12-06T20:32:41.368376Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Init: remove any existing configuration\n",
    "!rm -fr config\n",
    "!cp -r ../6-topical-rails/config .\n",
    "\n",
    "# Get rid of the TOKENIZERS_PARALLELISM warning\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f923f9cfe9e8f0f",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Prerequisites\n",
    "\n",
    "1. Install the `openai` package:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef8c379ded99a4db",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "!pip install openai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17f7d5ce578aaab8",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "2. Set the `OPENAI_API_KEY` environment variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "595f7001f160c3d6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-06T20:32:43.710660Z",
     "start_time": "2023-12-06T20:32:43.589636Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "!export OPENAI_API_KEY=$OPENAI_API_KEY    # Replace with your own key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0ab1d912ec76a6b",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "3. If you're running this inside a notebook, patch the AsyncIO loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b1181a203161cb75",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-06T20:50:14.514084Z",
     "start_time": "2023-12-06T20:50:14.502110Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fee3f3406f75ed6e",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Usage\n",
    "\n",
    "There are two modes in which you can use a guardrails configuration in conjunction with RAG:\n",
    "\n",
    "1. **Relevant Chunks**: perform the retrieval yourself and pass the **relevant chunks** directly to the `generate` method.\n",
    "2. **Knowledge Base**: configure a **knowledge base** directly into the guardrails configuration and let NeMo Guardrails manage the retrieval part.  \n",
    "\n",
    "### Relevant Chunks\n",
    "\n",
    "In the previous guide, the message \"How many free vacation days do I have per year\" yields a general response:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "116122bcb3caa890",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-06T20:50:29.935467Z",
     "start_time": "2023-12-06T20:50:17.142738Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Full-time employees are eligible for up to two weeks of paid vacation time per year. Part-time employees receive a prorated amount based on their hours worked. Please refer to the employee handbook for more information.\n"
     ]
    }
   ],
   "source": [
    "from nemoguardrails import LLMRails, RailsConfig\n",
    "\n",
    "config = RailsConfig.from_path(\"./config\")\n",
    "rails = LLMRails(config)\n",
    "\n",
    "response = rails.generate(messages=[{\"role\": \"user\", \"content\": \"How many vacation days do I have per year?\"}])\n",
    "print(response[\"content\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a1ccba02698781a",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "ABC company's Employee Handbook contains the following information:\n",
    "\n",
    "```markdown\n",
    "Employees are eligible for the following time off:\n",
    "\n",
    "* Vacation: 20 days per year, accrued monthly.\n",
    "* Sick leave: 15 days per year, accrued monthly.\n",
    "* Personal days: 5 days per year, accrued monthly.\n",
    "* Paid holidays: New Year's Day, Memorial Day, Independence Day, Thanksgiving Day, Christmas Day.\n",
    "* Bereavement leave: 3 days paid leave for immediate family members, 1 day for non-immediate family members.\n",
    "```\n",
    "\n",
    "You can pass this information directly to guardrails when making a `generate` call:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "28fce676db0c1900",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-06T20:50:40.534129Z",
     "start_time": "2023-12-06T20:50:34.593431Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Eligible employees receive 20 days of paid vacation time per year, which accrues monthly. You can find more information about this in the employee handbook.\n"
     ]
    }
   ],
   "source": [
    "response = rails.generate(\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"context\",\n",
    "            \"content\": {\n",
    "                \"relevant_chunks\": \"\"\"\n",
    "            Employees are eligible for the following time off:\n",
    "              * Vacation: 20 days per year, accrued monthly.\n",
    "              * Sick leave: 15 days per year, accrued monthly.\n",
    "              * Personal days: 5 days per year, accrued monthly.\n",
    "              * Paid holidays: New Year's Day, Memorial Day, Independence Day, Thanksgiving Day, Christmas Day.\n",
    "              * Bereavement leave: 3 days paid leave for immediate family members, 1 day for non-immediate family members. \"\"\"\n",
    "            },\n",
    "        },\n",
    "        {\"role\": \"user\", \"content\": \"How many vacation days do I have per year?\"},\n",
    "    ]\n",
    ")\n",
    "print(response[\"content\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b42b62f4fd791e3a",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "As expected, the response contains the correct answer. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5c09c2f83e25e33",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Knowledge Base\n",
    "\n",
    "There are three ways you can configure a knowledge base directly into a guardrails configuration: \n",
    "\n",
    "1. Using the *kb* folder.\n",
    "2. Using a custom `retrieve_relevant_chunks` action.\n",
    "3. Using a custom `EmbeddingSearchProvider`.\n",
    "\n",
    "For option 1, you can add a knowledge base directly into your guardrails configuration by creating a *kb* folder inside the *config* folder and adding documents there. Currently, only the Markdown format is supported. For a quick example, check out the complete implementation of the [ABC Bot](../../../examples/bots/abc).\n",
    "\n",
    "Options 2 and 3 represent advanced use cases beyond the scope of this topic."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7ba07763daafa2c",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Wrapping Up\n",
    "\n",
    "This guide introduced how a guardrails configuration can be used in the context of a RAG setup. \n",
    "\n",
    "## Next\n",
    "\n",
    "To continue learning about NeMo Guardrails, check out:\n",
    "1. [Guardrails Library](../../../docs/user-guides/guardrails-library.md).\n",
    "2. [Configuration Guide](../../../docs/user-guides/configuration-guide.md).\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
